from contextlib import contextmanager
from typing import Iterable

import torch
from transformers import PreTrainedModel


def determine_parameters(*,
    model: PreTrainedModel | None = None,
    modules: Iterable[torch.nn.Module | str] | None = None,
    parameters: Iterable[torch.Tensor | str] | None = None,
) -> frozenset[torch.Tensor]:
    if parameters is not None:
        parameters_ = []
        for parameter in parameters:
            if isinstance(parameter, torch.Tensor):
                parameters_.append(parameter)
                continue
            if isinstance(parameter, str):
                parameter = model.get_parameter(parameter)
                parameters_.append(parameter)
                continue
            raise TypeError(f"Unexpected parameter {parameter}!")
        return frozenset(parameters_)

    if modules is not None:
        parameters_ = []
        for module in modules:
            if isinstance(module, torch.nn.Module):
                parameters_.extend(module.parameters())
                continue
            if isinstance(module, str):
                module = model.get_submodule(module)
                parameters_.extend(module.parameters())
                continue
            raise TypeError(f"Unexpected module {module}!")
        return frozenset(parameters_)

    if model is not None:
        return frozenset(model.parameters())

    return frozenset()


@contextmanager
def restrict_requires_grad(
    model: PreTrainedModel, *,
    modules: Iterable[torch.nn.Module | str] | None = None,
    parameters: Iterable[torch.Tensor | str] | None = None,
):
    parameters = determine_parameters(
        model=model,
        modules=modules,
        parameters=parameters)

    changed = {}
    for parameter in model.parameters():
        requires_grad = parameter in parameters
        if parameter.requires_grad != requires_grad:
            changed[parameter] = parameter.requires_grad
            parameter.requires_grad = requires_grad

    try:
        yield
    finally:
        for parameter, org_requires_grad in changed.items():
            parameter.requires_grad = org_requires_grad
